//
// Created by YE on 2024/2/17.
//
#include "ych_mat.h"
#include "string.h"
#include "stdlib.h"
#include "math.h"
#include "stdio.h"
#include "../common/system.h"


void initMat(Mat *mat, int m, int n) {
    mat->rows = m, mat->cols = n, mat->is_popp = 0;
    mat->values = mallocMat(m, n);
}

void *createMat(int m, int n) {
    Mat *mat = malloc(sizeof(Mat));
    mat->rows = m, mat->cols = n, mat->is_popp = 0;
    mat->values = malloc(sizeof(Float *) * m);
    for (int i = m - 1; i >= 0; --i) {
        mat->values[i] = malloc(sizeof(Float) * n);
    }
    return mat;
}


void *copyMat(Mat *mat) {
    Mat *new_mat = malloc(sizeof(Mat));
    new_mat->rows = mat->rows, new_mat->cols = mat->cols, new_mat->is_popp = 0;
    new_mat->values = malloc(sizeof(Float *) * mat->rows);
    for (int i = mat->rows - 1; i >= 0; --i) {
        new_mat->values[i] = malloc(sizeof(Float) * mat->cols);
        memcpy(new_mat->values[i], mat->values[i], sizeof(Float) * mat->cols);
    }
    return new_mat;
}

Float **mallocMat(int m, int n) {
    Float **values = malloc(sizeof(Float *) * m);
    for (int i = m - 1; i >= 0; --i) {
        *(values + i) = malloc(sizeof(Float) * n);
    }
    return values;
}

void freeMat(Mat *mat) {
    for (int i = mat->rows - 1; i >= 0; --i) {
        free(mat->values[i]);
    }
    free(mat->values);
}

void deleteMat(Mat *mat) {
    for (int i = mat->rows - 1; i >= 0; --i) {
        free(mat->values[i]);
    }
    free(mat->values);
    free(mat);
}

static void freeMatValue(Float **values, int n) {
    for (int i = n - 1; i >= 0; --i) {
        free(values[i]);
    }
    free(values);
}

void matAdd(Float **matT, Float **matA, Float **matB, const int ROWS, const int COLS) {
    for (int i = ROWS - 1; i >= 0; --i)
        for (int j = COLS - 1; j >= 0; --j)
            matT[i][j] = matA[i][j] + matB[i][j];
}

void matSub(Float **matT, Float **matA, Float **matB, const int ROWS, const int COLS) {
    for (int i = ROWS - 1; i >= 0; --i)
        for (int j = COLS - 1; j >= 0; --j)
            matT[i][j] = matA[i][j] - matB[i][j];
}

void matMul(Float **matT, Float **matA, Float f, const int ROWS, const int COLS) {
    for (int i = ROWS - 1; i >= 0; --i)
        for (int j = COLS - 1; j >= 0; --j)
            matT[i][j] = matA[i][j] * f;
}

void matDiv(Float **matT, Float **matA, Float f, const int ROWS, const int COLS) {
    for (int i = ROWS - 1; i >= 0; --i)
        for (int j = COLS - 1; j >= 0; --j)
            matT[i][j] = matA[i][j] / f;
}

void matInner(Float **matT, Float **matA, Float **matB, const int ROWS, const int COLS,
              const int ICOLS) {
    for (int i = ROWS - 1; i >= 0; --i) {
        for (int j = COLS - 1; j >= 0; --j) {
            matT[i][j] = 0;
            for (int k = ICOLS - 1; k >= 0; --k) {
                matT[i][j] += matA[i][k] * matB[k][j];
            }
        }
    }
}

// 矩阵相乘的线程函数
void *imultiply(void *param) {
    MatParameters *params = (MatParameters *) param;
    int end_row = params->end_row;
    int end_col = params->end_col;
    int N = params->k;
    Float sum;
    Float **T = params->matT, **A = params->matA, **B = params->matB;
    for (int i = params->start_row; i < end_row; ++i) {
        for (int j = params->start_col; j < end_col; ++j) {
            sum = 0;
            for (int k = 0; k < N; ++k) {
                sum += A[i][k] * B[k][j];
            }
            T[i][j] = sum;
        }
    }
    pthread_exit(NULL);
}

void matInnerThread(Float **matT, Float **matA, Float **matB, const int ROWS, const int COLS,
                    const int ICOLS) {
    pthread_t *tid = malloc(sizeof(pthread_t) * max_threads);
    MatParameters *mat_params = malloc(sizeof(MatParameters) * max_threads);
    MatParameters *mp;
    int threads = 0;//记录线程数量
    int N = ROWS * COLS / max_threads + 1;//每个线程需要计算的元素个数，向上取整
    //每个线程分配到的行数和列数计算（边界另外考虑)
    double p = sqrt(1.0 * N / ROWS / COLS);
    int rpt = p * ROWS + 1, cpt = p * COLS + 1;
    //线程分配
    for (int row = 0; row < ROWS; row += rpt) {
        for (int col = 0; col < COLS; col += cpt) {
            mp = mat_params + threads;
            mp->matT = matT;
            mp->matA = matA;
            mp->matB = matB;
            mp->start_row = row;
            mp->end_row = row + rpt > ROWS ? ROWS : row + rpt;
            mp->start_col = col;
            mp->end_col = col + cpt > COLS ? COLS : col + cpt;
            mp->k = ICOLS;
            pthread_create(tid + threads, NULL, imultiply, mat_params + threads);
            threads++;
        }
    }
    for (int i = 0; i < threads; i++) {
        pthread_join(tid[i], NULL);
    }
    free(tid);
    free(mat_params);
}

int matEq(Float **matA, Float **matB, const int ROWS, const int COLS) {
    for (int i = ROWS - 1; i >= 0; --i) {
        for (int j = COLS - 1; j >= 0; --j) {
            if (matA[i][j] != matB[i][j])
                return 0;
        }
    }
    return 1;
}

void matPow(Float **matT, Float **matA, unsigned int n, const int N) {
    Float **base1 = mallocMat(N, N), **base2 = mallocMat(N, N), **resultT = mallocMat(N, N), **baseT;
    for (int i = N - 1; i >= 0; --i) {
        for (int j = N - 1; j >= 0; --j) {
            base1[i][j] = matA[i][j];
            if (i == j)
                resultT[i][j] = matT[i][j] = 1;
            else
                resultT[i][j] = matT[i][j] = 0;
        }
    }
    while (n > 0) {
        if (n & 1) {
            matInner(matT, resultT, base1, N, N, N);
        }
        for (int i = N - 1; i >= 0; --i) {
            for (int j = N - 1; j >= 0; --j) {
                resultT[i][j] = matT[i][j];
            }
        }
        matInner(base2, base1, base1, N, N, N);
        baseT = base1, base1 = base2, base2 = baseT;
        n >>= 1;
    }
    freeMatValue(base1, N), freeMatValue(base2, N), freeMatValue(resultT, N);
}